import random
import statistics
import time
from typing import List, Dict, Any

from ..database.connection import get_db_session
from ..repositories.hydrometric_repository import HydrometricRepository
from ..utils.logger import get_logger
from ..utils.time_utils import get_time_range_for_hours

# 配置日志
logger = get_logger(__name__)


class PerformanceTest:
    """数据库查询性能测试类"""

    def __init__(self):
        self.session = None
        self.repository = None
        self.test_results = []

    def __del__(self):
        if hasattr(self, 'db'):
            self.db.close()

    def get_random_station_ids(self, count: int = 10) -> List[str]:
        """获取随机测站ID列表"""
        with next(get_db_session()) as session:
            repository = HydrometricRepository(session)
            station_ids = repository.get_distinct_station_ids(1000)
            return random.sample(station_ids, min(count, len(station_ids)))

    def test_single_station_query(self, station_id: str, hours: int = 24) -> Dict[str, Any]:
        """测试单个测站查询性能"""
        start_time, end_time = get_time_range_for_hours(hours)

        results = {}

        with next(get_db_session()) as session:
            repository = HydrometricRepository(session)

            # 测试时序表查询
            start = time.time()
            timeseries_data = repository.query_by_station_and_time(
                station_id, start_time, end_time, 'timeseries'
            )
            timeseries_time = time.time() - start

            # 测试普通表查询
            start = time.time()
            regular_data = repository.query_by_station_and_time(
                station_id, start_time, end_time, 'regular'
            )
            regular_time = time.time() - start

            results = {
                'test_type': 'single_station_query',
                'station_id': station_id,
                'time_range_hours': hours,
                'timeseries_time': timeseries_time,
                'regular_time': regular_time,
                'timeseries_count': len(timeseries_data),
                'regular_count': len(regular_data),
                'performance_ratio': regular_time / timeseries_time if timeseries_time > 0 else 0
            }

        return results

    def test_multiple_stations_query(self, station_ids: List[str], hours: int = 24) -> Dict[str, Any]:
        """测试多个测站查询性能"""
        start_time, end_time = get_time_range_for_hours(hours)

        results = {}

        with next(get_db_session()) as session:
            repository = HydrometricRepository(session)

            # 测试时序表查询
            start = time.time()
            timeseries_data = repository.query_by_stations_and_time(
                station_ids, start_time, end_time, 'timeseries'
            )
            timeseries_time = time.time() - start

            # 测试普通表查询
            start = time.time()
            regular_data = repository.query_by_stations_and_time(
                station_ids, start_time, end_time, 'regular'
            )
            regular_time = time.time() - start

            results = {
                'test_type': 'multiple_stations_query',
                'station_count': len(station_ids),
                'time_range_hours': hours,
                'timeseries_time': timeseries_time,
                'regular_time': regular_time,
                'timeseries_count': len(timeseries_data),
                'regular_count': len(regular_data),
                'performance_ratio': regular_time / timeseries_time if timeseries_time > 0 else 0
            }

        return results

    def test_time_range_query(self, hours: int = 24) -> Dict[str, Any]:
        """测试时间范围查询性能"""
        start_time, end_time = get_time_range_for_hours(hours)

        results = {}

        with next(get_db_session()) as session:
            repository = HydrometricRepository(session)

            # 测试时序表查询
            start = time.time()
            timeseries_stats = repository.query_time_range_stats(
                start_time, end_time, 'timeseries'
            )
            timeseries_time = time.time() - start

            # 测试普通表查询
            start = time.time()
            regular_stats = repository.query_time_range_stats(
                start_time, end_time, 'regular'
            )
            regular_time = time.time() - start

            results = {
                'test_type': 'time_range_query',
                'time_range_hours': hours,
                'timeseries_time': timeseries_time,
                'regular_time': regular_time,
                'timeseries_stats': timeseries_stats,
                'regular_stats': regular_stats,
                'performance_ratio': regular_time / timeseries_time if timeseries_time > 0 else 0
            }

        return results

    def run_comprehensive_test(self, iterations: int = 5) -> List[Dict[str, Any]]:
        """运行综合性能测试"""
        logger.info(f"开始运行综合性能测试，迭代次数: {iterations}")

        all_results = []
        station_ids = self.get_random_station_ids(20)

        if not station_ids:
            logger.error("没有找到测站数据，请先运行数据生成脚本")
            return []

        for i in range(iterations):
            logger.info(f"执行第 {i + 1}/{iterations} 轮测试")

            # 测试1: 单站查询（不同时间范围）
            single_station_results = []
            for hours in [24, 168, 720]:  # 1天, 7天, 30天
                station_id = random.choice(station_ids)
                result = self.test_single_station_query(station_id, hours)
                single_station_results.append(result)

            # 测试2: 多站聚合查询
            multi_station_results = []
            for station_count in [5, 10, 20]:
                selected_stations = random.sample(station_ids, min(station_count, len(station_ids)))
                result = self.test_multiple_stations_query(selected_stations, 168)  # 7天
                multi_station_results.append(result)

            # 测试3: 时间范围查询
            time_range_results = []
            for hours in [24, 168, 720]:  # 1天, 7天, 30天
                result = self.test_time_range_query(hours)
                time_range_results.append(result)

            iteration_results = {
                'iteration': i + 1,
                'single_station': single_station_results,
                'multiple_stations': multi_station_results,
                'time_range': time_range_results
            }

            all_results.append(iteration_results)

        logger.info("综合性能测试完成")
        return all_results

    def analyze_results(self, results: List[Dict[str, Any]]) -> Dict[str, Any]:
        """分析测试结果"""
        if not results:
            return {}

        # 按测试类型分组
        grouped_results = {}
        for result in results:
            test_type = result.get('test_type', 'unknown')
            if test_type not in grouped_results:
                grouped_results[test_type] = []
            grouped_results[test_type].append(result)

        # 计算统计信息
        analysis = {}
        for test_type, type_results in grouped_results.items():
            timeseries_times = [r.get('timeseries_time', 0) for r in type_results]
            regular_times = [r.get('regular_time', 0) for r in type_results]
            ratios = [r.get('performance_ratio', 0) for r in type_results if r.get('performance_ratio', 0) > 0]

            analysis[test_type] = {
                'count': len(type_results),
                'timeseries_avg_time': statistics.mean(timeseries_times) if timeseries_times else 0,
                'regular_avg_time': statistics.mean(regular_times) if regular_times else 0,
                'avg_performance_ratio': statistics.mean(ratios) if ratios else 0,
                'timeseries_min_time': min(timeseries_times) if timeseries_times else 0,
                'timeseries_max_time': max(timeseries_times) if timeseries_times else 0,
                'regular_min_time': min(regular_times) if regular_times else 0,
                'regular_max_time': max(regular_times) if regular_times else 0
            }

        return analysis


def main():
    """主函数"""
    logger.info("开始性能测试")

    # 创建性能测试实例
    perf_test = PerformanceTest()

    try:
        # 运行综合测试
        results = perf_test.run_comprehensive_test(iterations=3)

        if not results:
            logger.error("测试失败：没有数据或连接问题")
            return 1

        # 分析结果
        analysis = perf_test.analyze_results(results)

        # 输出结果
        logger.info("=== 性能测试结果 ===")
        for test_type, stats in analysis.items():
            logger.info(f"\n{test_type} 测试结果:")
            logger.info(f"  TimescaleDB平均查询时间: {stats['timeseries_avg_time']:.4f}秒")
            logger.info(f"  普通表平均查询时间: {stats['regular_avg_time']:.4f}秒")
            logger.info(f"  性能提升比例: {stats['avg_performance_ratio']:.2f}x")
            logger.info(f"  测试次数: {stats['count']}")

    except Exception as e:
        logger.error(f"性能测试失败: {str(e)}")
        return 1

    return 0


if __name__ == "__main__":
    exit(main())
